Skip to content

Add swiglu + MXFP8 fused CuTe DSL kernel - #2

Open
wolfcomos wants to merge 1 commit into
mainfrom
agent/swiglu-mxfp8-unified
Open

Add swiglu + MXFP8 fused CuTe DSL kernel#2
wolfcomos wants to merge 1 commit into
mainfrom
agent/swiglu-mxfp8-unified

Conversation

@wolfcomos

@wolfcomos wolfcomos commented Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Owner

Summary

This PR adds one CuTe DSL kernel that fuses the swiglu with the RCEIL MXFP8 cast: gate/up are read once, the bf16 activation tensor is never materialized, and rowwise (1x32), colwise (32x1), or both scales come out
of a single pass. This kernel is bit-compatible with the standalone mxfp8 cute dsl quantizers. The consuming
composite autograd op lives in torchtitan branch https://github.com/wolfcomos/torchtitan/pull/2/changes (will refactor for upstream torchtitan PR soon), this torchao PR is kernels-only.

Performance

GB200, torch 2.14 nightly, CUDA 13.4, public DSL 4.7.0. Baseline is an already-fused-activation arm: torch.compile-fused SwiGLU (one static Triton kernel per shape) + the standalone CuTe DSL quantizers

Eager:

input_shape     direction    scales    baseline_us    fused_us  speedup    fused_gbps
--------------  -----------  --------  -------------  --------  ---------  ----------
(128, 128)      forward      rowwise          105.38     62.53  1.69x             1.3
(128, 128)      forward      colwise          115.47     65.70  1.76x             1.3
(128, 128)      forward      both             179.42     73.86  2.43x             1.3
(128, 128)      backward     rowwise          110.19     65.54  1.68x             2.0
(128, 128)      backward     colwise          116.06     66.93  1.73x             2.0
(128, 128)      backward     both             199.01     79.81  2.49x             2.1
(4096, 2048)    forward      rowwise          112.58     64.61  1.74x           653.3
(4096, 2048)    forward      colwise          113.76     65.79  1.73x           641.5
(4096, 2048)    forward      both             172.75     73.07  2.36x           696.0
(4096, 2048)    backward     rowwise          103.12     65.86  1.57x          1027.0
(4096, 2048)    backward     colwise          115.60     68.14  1.70x           992.5
(4096, 2048)    backward     both             176.45     75.60  2.33x          1123.5
(4096, 7168)    forward      rowwise          111.23     71.71  1.55x          2059.9
(4096, 7168)    forward      colwise          110.19     72.06  1.53x          2049.8
(4096, 7168)    forward      both             177.20     80.26  2.21x          2217.9
(4096, 7168)    backward     rowwise          118.11     72.06  1.64x          3284.8
(4096, 7168)    backward     colwise          120.19     75.10  1.60x          3151.8
(4096, 7168)    backward     both             181.25     80.22  2.26x          3705.5
(16384, 7168)   forward      rowwise          164.86    110.59  1.49x          5342.8
(16384, 7168)   forward      colwise          164.26    120.83  1.36x          4890.0
(16384, 7168)   forward      both             219.20    143.36  1.53x          4966.4
(16384, 7168)   backward     rowwise          406.53    155.65  2.61x          6083.4
(16384, 7168)   backward     colwise          407.55    196.61  2.07x          4816.0
(16384, 7168)   backward     both             514.43    219.14  2.35x          5426.2
(131072, 8192)  forward      rowwise         1372.16    904.19  1.52x          5974.7
(131072, 8192)  forward      colwise         1387.52   1022.98  1.36x          5280.9
(131072, 8192)  forward      both            1849.09   1216.08  1.52x          5352.9
(131072, 8192)  backward     rowwise         3676.03   1321.98  2.78x          6548.5
(131072, 8192)  backward     colwise         3708.94   1710.08  2.17x          5062.4
(131072, 8192)  backward     both            4624.83   1859.09  2.49x          5847.8

This host had application clocks pinned to 2062 Mhz.

End-to-end in torchtitan (DeepSeek-V3 16B, 4x GB200, EP=4, TP=1, batch 4, seq 4096, MXFP8 MoE recipe, 50 steps), three arms — the stock MXFP8 path, the fused-w13 MoE composite running the unfused activation + standalone quantizers, and the same composite backed by this PR's kernels:

arm                             tflops(11-50)  vs stock  median step ms  loss@50  peak GiB
------------------------------  -------------  --------  --------------  -------  --------
stock MXFP8                            381.05        --          777.84  2.78596     96.15
w13 composite, unfused quant           379.31    -0.46%          782.05  2.78251     99.47
w13 composite, fused (this PR)         390.64    +2.52%          763.35  2.78362     99.59

Net +2.5% throughput (+3.0% vs the structurally identical unfused composite the kernel-only effect), median step time -1.8%, loss parity (max spread 3.5e-3 at step 50), and the tightest step-time tail (max step 869 ms vs 1081 stock).

The same A/B under the BF16-wgrad recipe (mxfp8_rceil_wgrad_with_hp: the kernel emits rowwise-only casts, and the composite recomputes the activation in BF16 at the autograd boundary for the two wgrad GEMMs — a boundary with less redundant traffic for fusion to remove):

arm                             tflops(11-50)  vs unfused  median step ms  loss@50  peak GiB
------------------------------  -------------  ----------  --------------  -------  --------
w13 composite, unfused                 382.00          --          776.61  2.77950     99.47
w13 composite, fused (this PR)         386.48      +1.17%          771.34  2.78439     99.41

The fused kernel still wins here (+1.2% throughput, -0.7% median step, kernel-only), with the margin narrowed as expected for the cheaper boundary.

API

from torchao.prototype.moe_training.kernels.mxfp8 import (
    gated_act_mxfp8_forward,   # (gated_input [M, 2K] bf16, *, rowwise, colwise)
    gated_act_mxfp8_backward,  # (grad_h [M, K] bf16, gated_input [M, 2K] bf16, *, rowwise, colwise)
)
# both return (output_rowwise, output_colwise, scales_rowwise, scales_colwise)

gated_input packs [gate | up] — the layout a fused w13 projection produces; backward returns concatenated [dGate | dUp], so there is no torch.cat on the hot path. Disabled directions come back zero-sized (fixed arity). Registered as torch custom ops with fake impls; torch.compile(fullgraph=True) works. One kernel source: direction,
layout mode, and the activation (a module-level Constexpr policy function, _silu_pair) are compile-time parameters; shapes are runtime, so one compilation per mode serves every shape.

Numerics and tests

  • E8M0 scales and forward E4M3 data are bitwise identical to the eager activation followed by the standalone quantizers — including NaN/Inf, zero, subnormal-tiny, and E4M3-boundary amax regimes under the Refine MXFP8 CUDA quantization semantics pytorch/ao#4725 contract.
  • Backward E4M3 data may differ by at most one code in ≤1e-5 of elements (measured 5.7e-7): the kernel's PTX sigmoid and FMA contraction have no bit-exact eager equivalent without unfusing.
  • 80 cases, ~14 s on GB200: {fwd,bwd} × {rowwise,colwise,both} × edge-value patterns and a shape sweep; compile-fullgraph and cross-mode consistency checked bitwise; every launcher error path exercised (INT32 bound probed allocation-free via FakeTensorMode; pointer-alignment and geometry-override guards, each with a bitwise positive case); opt-in largest-legal-shape check against the standalone quantizer. Upstream's cutedsl test_mxfp8_dsl_special_value_semantics cases pass in the same environment.

Reproducing

pytest test/prototype/moe_training/test_cutedsl_gated_act_mxfp8.py -q

# MXFP8_BENCH_VALIDATE=1 checks every configuration against eager before timing;
# --shape/--direction/--scales run a single cell
MXFP8_BENCH_VALIDATE=1 python \
  benchmarks/prototype/moe_training/mxfp8/bench_gated_act_mxfp8_unified.py [--compile]

Follow-ups

  1. SiLU is the only activation exposed today, plan to add more activations in future PR.

  2. the TorchAO op itself has no register_sharding or automatic fallback. The external TorchTitan composite must choose the unfused path before invoking it, so there might be kernel sharding register work need to do for torchtitan integration.

@wolfcomos
wolfcomos force-pushed the agent/swiglu-mxfp8-unified branch from 9853226 to 26f6cf7 Compare August 9, 2026 08:10
@wolfcomos wolfcomos changed the title Add unified SwiGLU+MXFP8 CuTe DSL kernel with forward/backward and bi… Add gated ACT + MXFP8 fused CuTe DSL kernel with forward/backward and bi… Aug 11, 2026
@wolfcomos
wolfcomos force-pushed the agent/swiglu-mxfp8-unified branch 4 times, most recently from 22783a1 to 8ad8a79 Compare August 12, 2026 22:36
@wolfcomos wolfcomos changed the title Add gated ACT + MXFP8 fused CuTe DSL kernel with forward/backward and bi… Add swiglu + MXFP8 fused CuTe DSL kernel with forward/backward and bi… Aug 12, 2026
@wolfcomos wolfcomos changed the title Add swiglu + MXFP8 fused CuTe DSL kernel with forward/backward and bi… Add swiglu + MXFP8 fused CuTe DSL kernel Aug 12, 2026
@wolfcomos
wolfcomos force-pushed the agent/swiglu-mxfp8-unified branch 4 times, most recently from 486ceaf to 94489e8 Compare August 13, 2026 07:22
@wolfcomos
wolfcomos marked this pull request as ready for review August 13, 2026 07:25
@wolfcomos
wolfcomos force-pushed the agent/swiglu-mxfp8-unified branch 3 times, most recently from 39db529 to 61eb093 Compare August 19, 2026 05:14
@wolfcomos @claude
Add a fused gated-activation (SwiGLU) + MXFP8 quantization kernel for…
7e08b94 
… Blackwell

One CuTe DSL kernel fuses the gated activation (forward and backward)
with the RCEIL MXFP8 cast: gate/up are read once and the bf16 activation
never round-trips through global memory. Rowwise (1x32), colwise (32x1),
or both scale modes come out of a single pass, in the same blocked
tcgen05 layouts as the standalone quantizers, and bit-compatible with
their special-value semantics (pytorch#4725 contract: NaN/Inf amax invalidates
the block with scale byte 255 and all-NaN data; byte-0 scales descale by
2^127).

Ops: torchao::gated_act_mxfp8_{forward,backward} custom ops with fake
impls (torch.compile fullgraph works), public wrappers, 80-case numerics
suite (bitwise forward, one-code-bounded backward), and an A/B benchmark
vs the unfused Triton-activation + standalone-quantizer path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@wolfcomos
wolfcomos force-pushed the agent/swiglu-mxfp8-unified branch from 61eb093 to 7e08b94 Compare August 24, 2026 05:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@wolfcomos